package com.bruce.springai.repository;

import java.time.Instant;
import java.util.List;
import java.util.stream.Collectors;

import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.stereotype.Component;

import com.bruce.springai.bean.ChatMessage;

import jakarta.annotation.Resource;

/**
 * 使用MongoDB保存聊天记录
 */
@Component
public class MongoChatMemoryRepository implements ChatMemoryRepository{

    @Resource
    private MongoTemplate mongoTemplate ;

    @Override
    public List<String> findConversationIds() {
        return mongoTemplate.query(ChatMessage.class)
                .distinct("conversationId")
                .as(String.class)
                .all();
    }

    /**
     * 查找聊天记录
     */
    @Override
    public List<Message> findByConversationId(String conversationId) {
        List<ChatMessage> docs = mongoTemplate.find(
            Query.query(Criteria.where("conversationId").is(conversationId)), 
            ChatMessage.class
        );
        return docs.stream().map(ChatMessage::toMessage).collect(Collectors.toList());
    }

    /**
     * 保存聊天记录
     */
    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        Instant instant = Instant.now() ;
        List<ChatMessage> entities = messages.stream()
                .map(msg -> new ChatMessage(conversationId, msg,instant))
                .collect(Collectors.toList());
        mongoTemplate.insertAll(entities);
    }

    /**
     * 删除聊天记录
     */
    @Override
    public void deleteByConversationId(String conversationId) {
        mongoTemplate.remove(
            Query.query(Criteria.where("conversationId").is(conversationId))
        );
    }

}
